import os
import asyncio
from typing import Optional
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager

from ..config import settings
from ..models import ModelManager
from .endpoints import router


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize and shutdown logic for the FastAPI app."""
    # Startup
    app.state.model_manager = ModelManager()
    await app.state.model_manager.load_model(
        model_name=settings.model_name,
        backend=settings.backend,
        model_path=settings.model_path
    )
    
    yield
    
    # Shutdown
    if hasattr(app.state, 'model_manager'):
        await app.state.model_manager.unload_model()


def create_app() -> FastAPI:
    """Create and configure the FastAPI application."""
    app = FastAPI(
        title="BigServer API",
        description="A high-performance LLM serving framework supporting vLLM, SGLang, and Transformers",
        version="0.1.0",
        lifespan=lifespan
    )
    
    # Add CORS middleware
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    
    # Include API routes
    app.include_router(router, prefix="/v1", tags=["OpenAI-compatible API"])
    
    return app


def serve(host: str = "0.0.0.0", port: int = 8000, model_name: str = None, backend: str = None, model_path: str = None):
    """
    Start the BigServer API server.
    
    Args:
        host: Host address to bind the server to
        port: Port to run the server on
        model_name: Name of the model to load
        backend: Backend to use (vllm, sglang, transformers)
        model_path: Path to the model files
    """
    import uvicorn
    
    # Update settings if provided
    if model_name:
        settings.model_name = model_name
    if backend:
        settings.backend = backend
    if model_path:
        settings.model_path = model_path
    
    app = create_app()
    
    uvicorn.run(
        app,
        host=host,
        port=port,
        log_level="info"
    )


if __name__ == "__main__":
    serve()